#include "../kernel.h"

/* Store the final physical memory map */
static uint32_t   mlayout_count;
static physaddr_t mlayout_addr[MEMMAP_MAXCOUNT << 1];

static uint32_t        mmap_count;
static struct memmap_t mmap[MEMMAP_MAXCOUNT];

static uint32_t   mmap_node_count;
static physaddr_t mmap_node_addr[MEMMAP_MAXCOUNT << 1];
static int32_t    mmap_node_flag[MEMMAP_MAXCOUNT << 1];

#define PA_VAILD(addr) (mlayout_bin_search(addr) & 1)
/* return the first boundary index that has value larger than addr */
static int
mlayout_bin_search(physaddr_t addr)
{
     int l = 0, r = mlayout_count - 1, m;

     while (l < r)
     {
		  if (mlayout_addr[m = ((l + r) >> 1)] > addr)
			   r = m;
		  else l = m + 1;
     }

     return l;
}

static int          page_count;
static spin_lock_t *page_lock;
static int         *page_rc;

#define BUDDY_LEVEL_COUNT 20
/* Buddy system */
/* {{{ */
struct buddy_block_t
{
	 unsigned int type    : 2;
	 unsigned int status  : 2;
	 unsigned int ll_prev : 20;
	 unsigned int ll_next : 20;
	 unsigned int slot    : 4;
	 unsigned int pstatus : 16;
} __attribute__((packed));

#define SLOT_COUNT  2
#define SLOT_DMA    0
#define SLOT_NO_DMA 1

#define BBLOCK_ROOT      0
#define BBLOCK_LEFT      1
#define BBLOCK_RIGHT     2

#define BBLOCK_FREE      0
#define BBLOCK_PARTIAL   1
#define BBLOCK_USED      2

#define LEFT_CHILD(id, lv)   (id - (1 << ((lv) - 1)))
#define RIGHT_CHILD(id, lv)  (id + (1 << ((lv) - 1)))
#define LEFT_PARENT(id, lv)  (id + (1 << (lv)))
#define RIGHT_PARENT(id, lv) (id - (1 << (lv)))
#define LEFT_TAIL(id, lv)    (id - (1 << (lv)) + 1)
#define RIGHT_TAIL(id, lv)   (id + (1 << (lv)) - 1)

static int bblock_count;
static volatile struct buddy_block_t *bblocks;
static volatile int *ll_head[SLOT_COUNT];

static int bblock_max_count;
static int bleaf_max_count;
static int level_max_count;

static void
bblock_dump(int b, int l)
{
	 if (bblocks[b].status == BBLOCK_USED)
	 {
		  kprintf("[%08x, %08x]",
				  (LEFT_TAIL(b, l) >> 1) << PAGE_SHIFT,
				  ((RIGHT_TAIL(b, l) >> 1) << PAGE_SHIFT) + PAGE_SIZE - 1);
	 }
	 else if (bblocks[b].status == BBLOCK_PARTIAL)
	 {
		  bblock_dump(LEFT_CHILD(b, l), l - 1);
		  bblock_dump(RIGHT_CHILD(b, l), l - 1);
	 }
}

static void
add_to_bblock_level_list(int b, int l)
{
	 if (ll_head[bblocks[b].slot][l] == -1)
	 {
		  ll_head[bblocks[b].slot][l] = b;
		  bblocks[b].ll_prev =
			   bblocks[b].ll_next = b;
	 }
	 else
	 {
		  int next = ll_head[bblocks[b].slot][l];
		  int prev = bblocks[next].ll_prev;

		  bblocks[next].ll_prev = 
			   bblocks[prev].ll_next = b;
		  bblocks[b].ll_next = next;
		  bblocks[b].ll_prev = prev;
	 }
}

static void
remove_from_buddy_level_list(int b, int l)
{
	 if (bblocks[b].ll_prev == b)
	 {
		  ll_head[bblocks[b].slot][l] = -1;
	 }
	 else
	 {
		  int next = bblocks[b].ll_next;
		  int prev = bblocks[b].ll_prev;

		  bblocks[next].ll_prev = prev;
		  bblocks[prev].ll_next = next;

		  if (ll_head[bblocks[b].slot][l] == b)
			   ll_head[bblocks[b].slot][l] = next;
	 }
}

static int
build_bblock(int start, int level, int type, int slot)
{
	 if (level > 0)
	 {
		  build_bblock(start, level - 1, BBLOCK_LEFT, slot);
	 }

	 int result = bblock_count ++;
	 bblocks[result].slot = slot;
	 bblocks[result].type = type;
	 bblocks[result].ll_prev = bblocks[result].ll_next = result;
	 
	 if (level > 0)
	 {
		  build_bblock(start + (1 << (level - 1)), level - 1, BBLOCK_RIGHT, slot);
		  
		  int left = LEFT_CHILD(result, level);
		  int right = RIGHT_CHILD(result, level);

		  if (bblocks[left].status == BBLOCK_FREE &&
			  bblocks[right].status == BBLOCK_FREE)
		  {
			   bblocks[result].status = BBLOCK_FREE;
		  }
		  else
		  {
			   if (bblocks[left].status == BBLOCK_FREE)
			   {
					add_to_bblock_level_list(left, level - 1);
					bblocks[result].status = BBLOCK_PARTIAL;
			   }
			   else if (bblocks[right].status == BBLOCK_FREE)
			   {
					add_to_bblock_level_list(right, level - 1);
					bblocks[result].status = BBLOCK_PARTIAL;
			   }
			   else if (bblocks[left].status == BBLOCK_USED &&
						bblocks[right].status == BBLOCK_USED)
			   {
					bblocks[result].status = BBLOCK_USED;
			   }
			   else
			   {
					bblocks[result].status = BBLOCK_PARTIAL;
			   }
		  }
	 }
	 else
	 {
		  int a = mlayout_bin_search(start << PAGE_SHIFT);
		  int b = mlayout_bin_search(((start + 1) << PAGE_SHIFT) - 1);
		  if (a == b &&
			  mlayout_addr[a] > (start << PAGE_SHIFT) &&
			  (a & 1) == 1)
			   bblocks[result].status = BBLOCK_FREE;
		  else bblocks[result].status = BBLOCK_USED;
	 }


	 return result;
}

static int
build_zone(int start, int level, int slot)
{
	 int root = build_bblock(start, level, BBLOCK_ROOT, slot);
	 if (bblocks[root].status == BBLOCK_FREE)
		  add_to_bblock_level_list(root, level);
	 ++ bblock_count;
	 return root;
}

static void
build_buddy(int start, int end, int slot, int verbose)
{
	 int b;
	 int cur = start;
	 while (cur < end)
	 {
		  int level = bsr(end - cur);
		  b = build_zone(cur, level, slot);
		  if (verbose)
			   bblock_dump(b, level);
		  cur += 1 << level;
	 }
	 if (verbose)
		  kprintf("\n");
}

static int
bblock_alloc(int level, int slot)
{
	 int i;
	 for (i = level; i < level_max_count; ++ i)
	 {
		  if (ll_head[slot][i] != -1)
		  {
			   int cur = ll_head[slot][i];
			   remove_from_buddy_level_list(cur, i);
			   // Split free block
			   while (i > level)
			   {
					bblocks[cur].status = BBLOCK_PARTIAL;
					add_to_bblock_level_list(RIGHT_CHILD(cur, i), i - 1);
					cur = LEFT_CHILD(cur, i);
					-- i;
			   }
			   bblocks[cur].status = BBLOCK_USED;
			   return LEFT_TAIL(cur, i) >> 1;
		  }
	 }
	 return -1;
}

static int
bblock_free(int start)
{
	 int cur = start << 1;
	 int level = 0;
	 while (bblocks[cur].status != BBLOCK_USED)
	 {
		  if (bblocks[cur].type == BBLOCK_LEFT)
			   cur = LEFT_PARENT(cur, level);
		  else return -1;
		  ++ level;
	 }
	 // TODO: owner check?
	 bblocks[cur].status = BBLOCK_FREE;
	 // Merge free blocks
	 while (bblocks[cur].type != BBLOCK_ROOT)
	 {
		  int parent, sibling;
		  if (bblocks[cur].type == BBLOCK_LEFT)
		  {
			   parent = LEFT_PARENT(cur, level);
			   sibling = RIGHT_CHILD(parent, level + 1);
		  }
		  else
		  {
			   parent = RIGHT_PARENT(cur, level);
			   sibling = LEFT_CHILD(parent, level + 1);
		  }
		  
		  if (bblocks[sibling].status == BBLOCK_FREE)
		  {
			   remove_from_buddy_level_list(sibling, level);
			   bblocks[parent].status = BBLOCK_FREE;
		  }
		  else break;
		  cur = parent;
		  ++ level;
	 }
	 add_to_bblock_level_list(cur, level);
}

/* }}} */

/* Allocate memory space when initailizing system */
static physaddr_t init_free;

static void *
init_alloc(size_t size, size_t align, int verbose)
{
	 if (verbose)
		  kprintf("[init_alloc] init_free = %p, size = %d, align = 0x%x\n", init_free, size, align);

	 /* Find the first segment possible */
     int lid = mlayout_bin_search(init_free);
     lid |= 1;

     while (lid < mlayout_count)
     {
		  if (init_free < mlayout_addr[lid - 1])
			   init_free = mlayout_addr[lid - 1];

		  /* Find the continuous memory space that would contain the
		   * need */
		  if ((init_free |=  align - 1) + size < mlayout_addr[lid])
			   break;
		  else lid += 2;
     }

     /* NULL is failed */
     if (lid >= mlayout_count) return NULL;
     if (init_free + size >= KSIZE) return NULL;
     else ++ init_free;

     void *result = KADDR(init_free);
	 
     if (verbose)
		  kprintf("[init_alloc] result = %p\n", result);

     init_free += size;

	 return result;
}

static int
do_init_alloc(int verbose)
{
	 /* Alloc and init for buddy system */
	 /* {{{ */
	 if ((bblocks = (struct buddy_block_t *)
		  init_alloc(sizeof(struct buddy_block_t) * bblock_max_count, PAGE_SIZE, verbose))
		 == NULL) return -E_NO_MEM;
	 int i;
	 for (i = 0; i != SLOT_COUNT; ++i)
	 {
		  if ((ll_head[i] = (int *)
			   init_alloc(sizeof(int) * level_max_count, PAGE_SIZE, verbose))
			  == NULL) return -E_NO_MEM;
	 }
	 /* }}} */
	 page_lock = (spin_lock_t *)
		  init_alloc(sizeof(spin_lock_t) * page_count, PAGE_SIZE, verbose);
	 page_rc = (int *)
		  init_alloc(sizeof(int) * page_count, PAGE_SIZE, verbose);
	 /* For PROC Struct */
	 if ((procs = (struct proc_t*)
		  init_alloc(PROCESS_MAX_COUNT * sizeof(struct proc_t), PAGE_SIZE, verbose))
		 == NULL)
     	  return -E_NO_MEM;
	 /* For IRQ PIPE */
	 if ((irq_pipes = (kpipe_t *)
		  init_alloc(IRQ_COUNT * sizeof(kpipe_t), PAGE_SIZE, verbose)) == NULL)
		  return -E_NO_MEM;
	 /* For IPS and Services */
	 if ((ips_nodes = (struct ips_node_t *)
		  init_alloc((1 << (PROCESS_MAX_COUNT_SHIFT + PROCESS_IPS_COUNT_SHIFT)) *
					 sizeof(struct ips_node_t), PAGE_SIZE, verbose)) == NULL)
		  return -E_NO_MEM;
	 if ((ips_aps = (struct ips_ap_t *)
		  init_alloc(IPS_AP_COUNT *
					 sizeof(struct ips_ap_t), PAGE_SIZE, verbose)) == NULL)
		  return -E_NO_MEM;
	 return 0;
}

static int
pmem_make_layout(int verbose)
{
     uint32_t i, j, k, t;

     mmap_node_count = 0;

     for (i = 0; i != mmap_count; ++i)
     {
		  mmap_node_addr[mmap_node_count] = mmap[i].base;
		  mmap_node_flag[mmap_node_count] = mmap[i].flag;
		  ++mmap_node_count;

		  /* ignore if the end overflows */
		  if (!MEMMAP_ADDR_OVERFLOW(mmap[i].end))
		  {
			   mmap_node_addr[mmap_node_count] =  mmap[i].end;
			   mmap_node_flag[mmap_node_count] = -mmap[i].flag;
			   ++mmap_node_count;
		  }
     }

     /* Sort the mmap node table */
     for (i = 0; i != mmap_node_count; ++i)
     {
		  j = i;
		  for (k = i + 1; k != mmap_node_count; ++k)
		  {
			   if (mmap_node_addr[k] < mmap_node_addr[j])
					j = k;
		  }
		  t = mmap_node_addr[i]; mmap_node_addr[i] = mmap_node_addr[j]; mmap_node_addr[j] = t;
		  t = mmap_node_flag[i]; mmap_node_flag[i] = mmap_node_flag[j]; mmap_node_flag[j] = t;
     }

     /* for debugging */
     /* for (i = 0; i != mmapNodeCount; ++i) */
     /* 	  kprintf("addr : %x, flag : %x\n", mmapNodeAddr[i], mmapNodeFlag[i]); */

     /* Make memory layout table */
     int32_t vaild_level, reclaim_level;
     bool last_vaild, v;
     
     vaild_level = 0; reclaim_level = 0; v = false;
     mlayout_count = 0;

     i = 0;
     while (1)
     {
		  if (i >= mmap_node_count) break;

		  last_vaild = v;
		  while (1)
		  {
			   switch (mmap_node_flag[i])
			   {
			   case MEMMAP_FREE:
					++vaild_level;
					break;
			   case -MEMMAP_FREE:
					--vaild_level;
					break;
			   case MEMMAP_RECLAIMABLE:
					++reclaim_level;
					break;
			   case -MEMMAP_RECLAIMABLE:
					--reclaim_level;
					break;
			   default:
					if (mmap_node_flag[i] > 0)
						 --vaild_level;
					else ++vaild_level;
			   }

			   ++i;
			   if (i >= mmap_node_count) break;
			   if (mmap_node_addr[i - 1] != mmap_node_addr[i]) break;
		  }
	  
		  if (vaild_level == 0)
			   v = (reclaim_level > 0);
		  else v = (vaild_level > 0);

		  /* for debugging */
		  // kprintf("%x %d %d %d\n", mmapNodeAddr[i - 1], vaildLevel, reclaimLevel, v);

		  if (v != last_vaild)
		  {
			   mlayout_addr[mlayout_count] = mmap_node_addr[i - 1];
			   ++mlayout_count;
		  }
     }

	 if (verbose)
	 {
		  kprintf("Physical memory layout boundary:\n");
		  for (i = 0; i != mlayout_count; ++i)
		  {
			   kprintf(" [%d] = %x\n", i, mlayout_addr[i]);
		  }
	 }

	 page_count = mlayout_addr[mlayout_count - 1] >> PAGE_SHIFT;
	 return 0;
}

mcs_lock_t pmem_alloc_lock;

static int
pmem_allocator_init(int verbose)
{
	 mcs_init(&pmem_alloc_lock);
	 int i, j;
	 for (i = 0; i != page_count; ++ i)
		  spl_init(&page_lock[i]);

	 bblock_count = 0;
	 for (i = 0; i != SLOT_COUNT; ++ i)
	 {
 		  for (j = 0; j != level_max_count; ++ j)
			   ll_head[i][j] = -1;
	 }

	 if (page_count > DMA_BOUNDARY_PAGE)
	 {
		  build_buddy(0, DMA_BOUNDARY_PAGE, SLOT_DMA, verbose);
		  build_buddy(DMA_BOUNDARY_PAGE, page_count, SLOT_NO_DMA, verbose);
	 }
	 else
	 {
		  build_buddy(0, page_count, SLOT_DMA, verbose);
	 }
	 return 0;
}

int
pmem_init(void)
{
	 int i, err;

	 mmap_count = *(uint32_t *)MEMMAP_BASE;
     memmove(mmap, MEMMAP_ENTRY, MEMMAP_SIZE * mmap_count);

     if (mmap_count == 0)
     {
		  kprintf("Cannot detect physical memory map.\n");
		  return -E_FAULT;
     }

     for (i = 0; i != mmap_count; ++i)
		  mmap[i].end += mmap[i].base;

	 /* mark the lower memory reserved */
     mmap[mmap_count].base = (uintptr_t)0;
     mmap[mmap_count].end =  (uintptr_t)0x00100000;
     mmap[mmap_count].flag = MEMMAP_RESERVED;
	 ++ mmap_count;

     /* generate the initial layout */
     pmem_make_layout(0);

	 bleaf_max_count = page_count;
	 bblock_max_count = bleaf_max_count << 1;
	 level_max_count = BUDDY_LEVEL_COUNT;

	 physaddr_t init_alloc_start;
	 init_alloc_start = DMA_BOUNDARY_PAGE << PAGE_SHIFT;
	 if (PADDR(_end) > init_alloc_start)
		  init_alloc_start = PADDR(_end);
	 init_free = init_alloc_start;

	 mmap[mmap_count].base = (uintptr_t)PADDR(_text);
	 mmap[mmap_count].end  = (uintptr_t)PADDR(_end);
	 mmap[mmap_count].flag = MEMMAP_RESERVED;
	 ++ mmap_count;

	 if ((err = do_init_alloc(0)) < 0)
		  return err;
	 
	 mmap[mmap_count].base = init_alloc_start;
	 mmap[mmap_count].end  = init_free;
	 mmap[mmap_count].flag = MEMMAP_RESERVED;
	 ++ mmap_count;

	 pmem_make_layout(0);

	 /* build the buddy system */
	 pmem_allocator_init(0);
	 return 0;
}

int
pmem_page_alloc(int num, physaddr_t *addr, int dma, int status)
{
	 if (status != PPAGE_MANAGED)
		  if (num > 1) return -E_INVAL;
	 
	 int level = bsr(num);
	 if (num > (1 << level)) ++ level;
	 
	 struct mcs_lock_node_t node;
	 mcs_acquire(&pmem_alloc_lock, &node);
	 int result = -1;
	 if (dma == 0 && page_count > DMA_BOUNDARY_PAGE)
		  result = bblock_alloc(level, SLOT_NO_DMA);
	 if (result == -1)
		  result = bblock_alloc(level, SLOT_DMA);
	 mcs_release(&pmem_alloc_lock, &node);
	 if (result == -1)
	 {
		  // kprintf("pmem page alloc failed wieh dma = %d num = %d\n", dma, num);
		  return -E_NO_MEM;
	 }
	 else
	 {
		  *addr = ((physaddr_t)result) << PAGE_SHIFT;
		  int i;
		  for (i = result; i != result + num; ++ i)
			   bblocks[i << 1].pstatus = status;
		  page_rc[result] = 0;
		  return 0;
	 }
}

void
pmem_page_free(physaddr_t addr)
{
	 struct mcs_lock_node_t node;
	 int page = addr >> PAGE_SHIFT;
	 if (page >= page_count) return;
	 mcs_acquire(&pmem_alloc_lock, &node);
	 bblock_free(page);
	 mcs_release(&pmem_alloc_lock, &node);
}

void
pmem_page_add_ref(physaddr_t addr, int delta)
{
	 int page = addr >> PAGE_SHIFT;
	 if (page >= page_count || bblocks[page << 1].status != PPAGE_CORE) return;
	 spl_acquire(&page_lock[page]);
	 page_rc[page] += delta;
	 spl_release(&page_lock[page]);
}

void
pmem_page_inc_ref(physaddr_t addr)
{
	 int page = addr >> PAGE_SHIFT;
	 if (page >= page_count || bblocks[page << 1].status != PPAGE_CORE) return;
	 spl_acquire(&page_lock[page]);
	 ++ page_rc[page];
	 spl_release(&page_lock[page]);
}

void
pmem_page_dec_ref(physaddr_t addr)
{
	 int page = addr >> PAGE_SHIFT;
	 if (page >= page_count || bblocks[page << 1].status != PPAGE_CORE) return;
	 spl_acquire(&page_lock[page]);
	 if (-- page_rc[page] == 0)
	 {
		  spl_release(&page_lock[page]);
		  bblock_free(page);
	 }
	 else spl_release(&page_lock[page]);
}

int
pmem_page_status(physaddr_t addr)
{
	 int result = 0;
	 int page = addr >> PAGE_SHIFT;
	 if (page >= page_count) return -1;
	 result = bblocks[page << 1].pstatus;
	 return result;
}

int
pmem_page_fsm(physaddr_t addr, int *tf_vec, int size)
{
	 int result = 0;
	 int page = addr >> PAGE_SHIFT;
	 if (page >= page_count) return -1;
	 spl_acquire(&page_lock[page]);
	 if (bblocks[page << 1].pstatus < size &&
		 tf_vec[bblocks[page << 1].pstatus] != -1)
		  result = bblocks[page << 1].pstatus = tf_vec[bblocks[page << 1].pstatus];
	 else result = -1;
	 spl_release(&page_lock[page]);
	 return result;
}
